from feature_extract import myutils
from feature_extract.registry import MODEL_REGISTRY
import torch
from torch import nn


@MODEL_REGISTRY.register('base_load')
def base_load(model: nn.Module, path, device=torch.device('cpu'), strict=True, **kwargs):
    state_dict = torch.load(path, map_location=device)
    model.load_state_dict(state_dict, strict=strict)
    return model


@MODEL_REGISTRY.register('extract_load')
def extract_load(model: nn.Module, path, device=torch.device('cpu'), dict_name='state_dict', strict=True, **kwargs):
    state_dict = torch.load(path, map_location=device)
    state_dict = state_dict[dict_name]
    model.load_state_dict(state_dict, strict=strict)
    return model


@MODEL_REGISTRY.register('adapted_extract_load')
def adapted_extract_load(model: nn.Module, path, device=torch.device('cpu'), dict_name='state_dict', strict=True,
                         **kwargs):
    state_dict = torch.load(path, map_location=device, encoding='iso-8859-1')
    state_dict = state_dict[dict_name]
    model.load_state_dict(state_dict, strict=strict)
    return model


@MODEL_REGISTRY.register('remove_prefix_extract_load')
def remove_prefix_extract_load(model: nn.Module, path, device=torch.device('cpu'), prefix="module.",
                               dict_name='state_dict', strict=True, **kwargs):
    state_dict = torch.load(path, map_location=device, encoding='iso-8859-1',weights_only=False)
    state_dict = state_dict[dict_name]
    state_dict = myutils.remove_prefix(state_dict, prefix)
    model.load_state_dict(state_dict, strict=strict)
    return model


def remove_add_prefix_extract_load(model: nn.Module, path, device=torch.device('cpu'), prefix="module.",
                               dict_name='state_dict', strict=True,prefix_add=None, **kwargs):
    state_dict = torch.load(path, map_location=device, encoding='iso-8859-1')
    state_dict = state_dict[dict_name]
    state_dict = myutils.remove_prefix(state_dict, prefix)
    if prefix_add is not None:
        state_dict=myutils.add_prefix(state_dict,prefix_add)
    model.load_state_dict(state_dict, strict=strict)
    return model

@MODEL_REGISTRY.register('remove_prefix_load')
def remove_prefix_load(model: nn.Module, path, device=torch.device('cpu'), prefix="module.", strict=True, **kwargs):
    state_dict = torch.load(path, map_location=device, encoding='iso-8859-1')
    state_dict = myutils.remove_prefix(state_dict, prefix)
    model.load_state_dict(state_dict, strict=strict)
    return model
